"""
MCP 使用示例
演示如何使用 MCP 客户端进行数据查询和分析
"""

import asyncio
import logging
import yaml
from typing import Dict, List

from mcp_rag.mcp.client.client import MCPClient

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

async def load_config() -> Dict:
    """加载配置文件"""
    with open("config/mcp_config.yaml", "r") as f:
        config = yaml.safe_load(f)
    return config["client"]

async def run_query_example(client: MCPClient):
    """查询示例"""
    # 执行简单查询
    results = await client.query(
        query="查找最近30天的销售数据",
        filters={"category": "电子产品"},
        limit=5
    )
    
    if results:
        logger.info("查询结果:")
        for item in results:
            logger.info(f"- {item}")
    else:
        logger.warning("未找到匹配的数据")
        
async def run_analysis_example(client: MCPClient):
    """分析示例"""
    # 准备分析数据
    data = [
        {"date": "2024-03-01", "sales": 1000, "category": "电子产品"},
        {"date": "2024-03-02", "sales": 1500, "category": "电子产品"},
        {"date": "2024-03-03", "sales": 800, "category": "电子产品"}
    ]
    
    # 执行趋势分析
    results = await client.analyze(
        data=data,
        analysis_type="trend",
        params={
            "metrics": ["sales"],
            "dimensions": ["date"],
            "window_size": 3
        }
    )
    
    if results:
        logger.info("分析结果:")
        logger.info(f"趋势: {results.get('trend')}")
        logger.info(f"预测: {results.get('forecast')}")
    else:
        logger.warning("分析失败")
        
async def run_stream_example(client: MCPClient):
    """流处理示例"""
    # 启动数据流
    stream_id = "sales_stream_1"
    result = await client.stream(
        stream_id=stream_id,
        action="start",
        params={
            "source": "sales_events",
            "filters": ["category=电子产品"],
            "window_size": 60  # 60秒窗口
        }
    )
    
    if result:
        logger.info(f"流处理已启动: {stream_id}")
        
        # 模拟接收数据
        for _ in range(3):
            result = await client.stream(
                stream_id=stream_id,
                action="receive"
            )
            if result:
                logger.info(f"收到数据: {result}")
            await asyncio.sleep(1)
            
        # 停止数据流
        await client.stream(
            stream_id=stream_id,
            action="stop"
        )
        logger.info("流处理已停止")
    else:
        logger.warning("流处理启动失败")
        
async def run_update_example(client: MCPClient):
    """更新示例"""
    # 更新产品信息
    success = await client.update(
        entity_id="product_123",
        updates={
            "price": 999.99,
            "stock": 100,
            "last_updated": "2024-03-20T10:00:00Z"
        }
    )
    
    if success:
        logger.info("更新成功")
    else:
        logger.warning("更新失败")
        
async def main():
    """主函数"""
    try:
        # 加载配置
        config = await load_config()
        
        # 创建客户端
        client = MCPClient(config)
        
        # 连接服务器
        if not await client.connect():
            logger.error("无法连接到服务器")
            return
            
        try:
            # 运行示例
            await run_query_example(client)
            await run_analysis_example(client)
            await run_stream_example(client)
            await run_update_example(client)
            
        finally:
            # 断开连接
            await client.disconnect()
            
    except Exception as e:
        logger.error(f"发生错误: {e}")
        
if __name__ == "__main__":
    # 运行示例
    asyncio.run(main()) 